package com.demo.wrapper;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.map.MapUtil;
import com.demo.wrapper.enums.BoolEnum;
import com.demo.wrapper.enums.OperatorEnum;
import com.demo.wrapper.enums.OrderEnum;
import com.demo.wrapper.enums.ScoreModeEnum;
import com.demo.wrapper.pojo.Agg;
import com.demo.wrapper.pojo.Arg;
import com.demo.wrapper.pojo.ConditionalReceiver;
import com.demo.wrapper.pojo.Order;
import com.demo.wrapper.utils.ArrayUtil;
import com.demo.wrapper.utils.FieldValueUtil;
import org.apache.lucene.search.join.ScoreMode;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.Operator;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.BucketOrder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;

import java.util.*;
import java.util.function.Consumer;

/**
 * es条件拼接器处理
 *
 * @author 19242
 */
public class WrapperUtil {

    public static <T extends AbstractEsQueryWrapper<T, ?, QueryBuilder, SearchSourceBuilder>> QueryBuilder getBuilder(T wrapper, boolean autoFilter) {
        if (CollectionUtil.isEmpty(wrapper.getChildren(false))) {
            Set<Arg> querySet = wrapper.getQuerySet(false);
            if (CollectionUtil.isEmpty(querySet)) {
                return null;
            }
            // 如果没有子查询且只有一条查询条件，是and或or关系的，返回单查询
            if (autoFilter && 1 == querySet.size()) {
                if (BoolEnum.MUST == wrapper.getBoolEnum() || BoolEnum.SHOULD == wrapper.getBoolEnum()) {
                    return WrapperUtil.setQueryBuilder(querySet.stream().findFirst().orElse(null), true);
                }
            }
        }
        return getBuilder(wrapper, autoFilter, null);
    }

    /**
     * 将wrapper收集到的条件转换成QueryBuilder
     *
     * @param wrapper       解析的拼接器
     * @param autoFilter    自动过滤相同关系的子条件，不支持not
     * @param parentBuilder 父节点的QueryBuilder
     * @param <T>           泛型
     * @return QueryBuilder
     */
    static <T extends AbstractEsQueryWrapper<T, ?, QueryBuilder, SearchSourceBuilder>> QueryBuilder getBuilder(T wrapper, boolean autoFilter,
                                                                                                               BoolQueryBuilder parentBuilder) {
        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
        Map<BoolEnum, List<T>> children = wrapper.getChildren(false);
        if (MapUtil.isNotEmpty(children)) {
            // not的不支持过滤
            if (autoFilter && wrapper.getBoolEnum() != BoolEnum.MUST_NOT) {
                // 当前节点和子节点是相同关系，把子节点条件并入当前条件
                List<T> common = children.remove(wrapper.getBoolEnum());
                mergeCommonChildren(wrapper, common);
            }
            // 不同关系的子节点收集起来
            for (Map.Entry<BoolEnum, List<T>> entry : children.entrySet()) {
                for (T esQueryWrapper : entry.getValue()) {
                    setBoolBuilder(boolQueryBuilder, esQueryWrapper.getBoolEnum(),
                            getBuilder(esQueryWrapper, autoFilter, boolQueryBuilder));
                }
            }
        }
        if (CollectionUtil.isNotEmpty(wrapper.getQuerySet(false))) {
            // 优先使用父节点拼接条件
            BoolQueryBuilder builder = Objects.nonNull(parentBuilder) ? parentBuilder : boolQueryBuilder;
            for (Arg arg : wrapper.getQuerySet(false)) {
                QueryBuilder queryBuilder = setQueryBuilder(arg, autoFilter);
                if (Objects.isNull(queryBuilder)) {
                    continue;
                }
                setBoolBuilder(builder, wrapper.getBoolEnum(), queryBuilder);
            }
        }
        return nonEmpty(boolQueryBuilder);
    }

    /**
     * 条件为空 返回null
     *
     * @param boolQueryBuilder
     * @return
     */
    private static QueryBuilder nonEmpty(BoolQueryBuilder boolQueryBuilder) {
        if (Objects.isNull(boolQueryBuilder)) {
            return null;
        }
        if (CollectionUtil.isEmpty(boolQueryBuilder.must()) && CollectionUtil.isEmpty(boolQueryBuilder.should())
                && CollectionUtil.isEmpty(boolQueryBuilder.mustNot()) && CollectionUtil.isEmpty(boolQueryBuilder.filter())) {
            return null;
        }
        return boolQueryBuilder;
    }

    /**
     * 合并相同关系的子查询
     *
     * @param wrapper 当前节点
     * @param common  关系相同的子节点
     */
    private static <Wrapper extends AbstractEsQueryWrapper<Wrapper, ?, QueryBuilder, SearchSourceBuilder>> void mergeCommonChildren(Wrapper wrapper,
                                                                                                                                    List<Wrapper> common) {
        if (CollectionUtil.isEmpty(common)) {
            return;
        }
        for (Wrapper esQueryWrapper : common) {
            // 将子节点的条件放入当前节点
            wrapper.addAllQuery(esQueryWrapper.getQuerySet(false));
            Map<BoolEnum, List<Wrapper>> children = esQueryWrapper.getChildren(false);
            if (CollectionUtil.isNotEmpty(children)) {
                List<Wrapper> remove = children.remove(wrapper.getBoolEnum());
                // 遍历子节点的子节点添加相同节点
                mergeCommonChildren(wrapper, remove);
                // 清空子节点的查询条件
                esQueryWrapper.clearQuery();
                // 将子查询放回
                wrapper.addChildren(esQueryWrapper);
            }
        }
    }

    /**
     * 添加关系
     *
     * @param boolQueryBuilder 多条件拼接
     * @param boolEnum         关系类型
     * @param queryBuilder     条件
     */
    static void setBoolBuilder(BoolQueryBuilder boolQueryBuilder, BoolEnum boolEnum, QueryBuilder queryBuilder) {
        if (Objects.isNull(queryBuilder)) {
            return;
        }
        switch (boolEnum) {
            case MUST:
                boolQueryBuilder.must(queryBuilder);
                break;
            case SHOULD:
                boolQueryBuilder.should(queryBuilder);
                break;
            case MUST_NOT:
                boolQueryBuilder.mustNot(queryBuilder);
                break;
            case FILTER:
                boolQueryBuilder.filter(queryBuilder);
            default:
                break;
        }
    }

    /**
     * 拼接查询条件
     *
     * @param arg        查询参数
     * @param autoFilter
     * @return 条件
     */
    static QueryBuilder setQueryBuilder(Arg arg, boolean autoFilter) {
        if (Objects.isNull(arg) || Objects.isNull(arg.getQueryEnum())) {
            return null;
        }
        switch (arg.getQueryEnum()) {
            case TERM:
                return QueryBuilders.termQuery(arg.getColumn(), arg.getVal()).boost(arg.getBoost());
            case TERMS:
                return QueryBuilders.termsQuery(arg.getColumn(), (Collection<?>) arg.getVal()).boost(arg.getBoost());
            case MATCH:
                return QueryBuilders.matchQuery(arg.getColumn(), arg.getVal()).boost(arg.getBoost()).operator(Operator.fromString(((OperatorEnum) arg.getExt1()).name()));
            case NESTED:
                return QueryBuilders.nestedQuery((String) arg.getExt1(), (QueryBuilder) ((AbstractEsQueryWrapper<?, ?, ?, ?>) arg.getExt2()).queryBuilder(autoFilter), getScoreMode((ScoreModeEnum) arg.getExt3()));
            case MATCH_ALL:
                return QueryBuilders.matchAllQuery().boost(arg.getBoost());
            case MATCH_PHRASE:
                return QueryBuilders.matchPhraseQuery(arg.getColumn(), arg.getVal()).boost(arg.getBoost());
            case MATCH_PHRASE_PREFIX:
                return QueryBuilders.matchPhrasePrefixQuery(arg.getColumn(), arg.getVal()).boost(arg.getBoost()).maxExpansions((int) arg.getExt1());
            case MULTI_MATCH:
                return QueryBuilders.multiMatchQuery(arg.getVal(), (String[]) arg.getExt1());
            case GT:
                return QueryBuilders.rangeQuery(arg.getColumn()).gt(arg.getVal()).boost(arg.getBoost());
            case GE:
                return QueryBuilders.rangeQuery(arg.getColumn()).gte(arg.getVal()).boost(arg.getBoost());
            case LT:
                return QueryBuilders.rangeQuery(arg.getColumn()).lt(arg.getVal()).boost(arg.getBoost());
            case LE:
                return QueryBuilders.rangeQuery(arg.getColumn()).lte(arg.getVal()).boost(arg.getBoost());
            case BETWEEN:
                return QueryBuilders.rangeQuery(arg.getColumn()).gte(arg.getExt1()).lte(arg.getExt2()).boost(arg.getBoost());
            case WILDCARD:
                return QueryBuilders.wildcardQuery(arg.getColumn(), arg.getVal().toString()).boost(arg.getBoost());
            case EXISTS:
                return QueryBuilders.existsQuery(arg.getColumn()).boost(arg.getBoost());
            case QUERY_BUILDER:
                return (QueryBuilder) arg.getExt1();
            default:
                return null;
        }
    }

    /**
     * 获取嵌套聚合得分模式
     *
     * @param scoreModeEnum
     * @return
     */
    private static ScoreMode getScoreMode(ScoreModeEnum scoreModeEnum) {
        if (Objects.isNull(scoreModeEnum)) {
            return ScoreMode.None;
        }
        switch (scoreModeEnum) {
            case Avg:
                return ScoreMode.Avg;
            case Max:
                return ScoreMode.Max;
            case Min:
                return ScoreMode.Min;
            case Total:
                return ScoreMode.Total;
            default:
                return ScoreMode.None;
        }
    }

    /**
     * 获取SearchSourceBuilder
     *
     * @param <Wrapper>  AbstractEsQueryWrapper的子类
     * @param wrapper    条件
     * @param autoFilter
     * @return
     */
    static <Wrapper extends AbstractEsQueryWrapper<Wrapper, ?, QueryBuilder, SearchSourceBuilder>> SearchSourceBuilder getSourceBuilder(Wrapper wrapper, boolean autoFilter) {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(wrapper.queryBuilder(autoFilter));
        List<Arg> limitList = wrapper.getLimitList(false);
        if (CollectionUtil.isNotEmpty(limitList)) {
            // 防止有些条件被覆盖掉，先收集完再处理
            ConditionalReceiver receiver = new ConditionalReceiver();
            for (Arg arg : limitList) {
                setLimit(searchSourceBuilder, arg, receiver);
            }
            String[] includes = receiver.getIncludes().toArray(new String[0]);
            String[] excludes = receiver.getExcludes().toArray(new String[0]);
            if (includes.length > 0 || excludes.length > 0) {
                searchSourceBuilder.fetchSource(includes, excludes);
            }
        }
        return searchSourceBuilder;
    }

    /**
     * 添加条件
     *
     * @param searchSourceBuilder
     * @param arg
     * @return
     */
    private static SearchSourceBuilder setLimit(SearchSourceBuilder searchSourceBuilder, Arg arg, ConditionalReceiver receiver) {
        if (Objects.isNull(arg) || Objects.isNull(arg.getQueryEnum())) {
            return searchSourceBuilder;
        }
        switch (arg.getQueryEnum()) {
            case FROM:
                return searchSourceBuilder.from((int) arg.getExt1());
            case SIZE:
                return searchSourceBuilder.size((int) arg.getExt1());
            case LIMIT:
                return searchSourceBuilder.from((int) arg.getExt1()).size((int) arg.getExt2());
            case INCLUDES:
                ArrayUtil.addAll(receiver.getIncludes(), (String) arg.getExt1(), (String[]) arg.getExt2());
                return searchSourceBuilder;
            case EXCLUDES:
                ArrayUtil.addAll(receiver.getExcludes(), (String) arg.getExt1(), (String[]) arg.getExt2());
                return searchSourceBuilder;
            case FETCH_SOURCE:
                ArrayUtil.addAll(receiver.getIncludes(), (String[]) arg.getExt1());
                ArrayUtil.addAll(receiver.getExcludes(), (String[]) arg.getExt2());
                return searchSourceBuilder;
            case AGG:
                return searchSourceBuilder.aggregation((AggregationBuilder) arg.getExt1());
            case SORT:
                return searchSourceBuilder.sort(arg.getColumn(), FieldValueUtil.sortOrder((OrderEnum) arg.getExt1()));
            case SOURCE_BUILDER:
                ((Consumer<SearchSourceBuilder>) arg.getExt1()).accept(searchSourceBuilder);
                return searchSourceBuilder;
            default:
                return searchSourceBuilder;
        }
    }

    /**
     * 展示字段设置
     *
     * @param searchSourceBuilder searchSourceBuilder
     * @param includes            展示字段
     * @param excludes            不展示字段
     * @return
     */
    private static SearchSourceBuilder fetchSource(SearchSourceBuilder searchSourceBuilder, String[] includes, String[] excludes) {
        if (Objects.isNull(includes) || includes.length == 0) {
            if (Objects.isNull(excludes) || excludes.length == 0) {
                return searchSourceBuilder;
            }
        }
        FetchSourceContext fetchSourceContext = searchSourceBuilder.fetchSource();
        if (Objects.isNull(fetchSourceContext)) {
            return searchSourceBuilder.fetchSource(includes, excludes);
        }
        // 合并旧数组和新数组
        String[] newIncludes = ArrayUtil.mergeArr(includes, fetchSourceContext.includes());
        String[] newExcludes = ArrayUtil.mergeArr(excludes, fetchSourceContext.excludes());

        return searchSourceBuilder.fetchSource(newIncludes, newExcludes);
    }


    /**
     * 组装聚合条件
     *
     * @param agg
     * @return
     */
    public static AggregationBuilder setAgg(Agg agg) {
        if (Objects.isNull(agg) || Objects.isNull(agg.getAggEnum())) {
            return null;
        }
        switch (agg.getAggEnum()) {
            case TERMS:
                Order[] orders = (Order[]) agg.getExt2();
                List<BucketOrder> bucketOrders = new ArrayList<>();
                for (Order order : orders) {
                    bucketOrders.add(BucketOrder.aggregation(order.getColumn(), order.isAsc()));
                }
                return AggregationBuilders.terms(agg.getName()).field(agg.getColumn()).size((int) agg.getExt1()).order(bucketOrders);
            case VALUE_COUNT:
                return AggregationBuilders.count(agg.getName()).field(agg.getColumn());
            case CARDINALITY:
                return AggregationBuilders.cardinality(agg.getName()).field(agg.getColumn());
            case AVG:
                return AggregationBuilders.avg(agg.getName()).field(agg.getColumn());
            case MAX:
                return AggregationBuilders.max(agg.getName()).field(agg.getColumn());
            case MIN:
                return AggregationBuilders.min(agg.getName()).field(agg.getColumn());
            case SUM:
                return AggregationBuilders.sum(agg.getName()).field(agg.getColumn());
            case STATS:
                return AggregationBuilders.stats(agg.getName()).field(agg.getColumn());
            case BUILDER:
                return (AggregationBuilder) agg.getExt1();
            case FILTER:
                return AggregationBuilders.filter(agg.getName(), (QueryBuilder) ((AbstractEsQueryWrapper<?, ?, ?, ?>) agg.getExt1()).queryBuilder());
            default:
                return null;
        }
    }

    /**
     * 将收集的Agg转成AggregationBuilder
     *
     * @return AggregationBuilder
     */
    public static AggregationBuilder getAggregationBuilder(AbstractEsAggWrapper<? extends AbstractEsAggWrapper<?, ?>, ?> aggWrapper) {
        AggregationBuilder aggregationBuilder = WrapperUtil.setAgg(aggWrapper.getAgg());
        if (Objects.isNull(aggregationBuilder)) {
            return null;
        }
        List<? extends AbstractEsAggWrapper<? extends AbstractEsAggWrapper<?, ?>, ?>> aggChildren = aggWrapper.getAggChildren(false);
        if (CollectionUtil.isEmpty(aggChildren)) {
            return aggregationBuilder;
        }
        // 遍历子聚合
        for (AbstractEsAggWrapper<? extends AbstractEsAggWrapper<?, ?>, ?> aggChild : aggChildren) {
            Optional.ofNullable(getAggregationBuilder(aggChild)).ifPresent(aggregationBuilder::subAggregation);
        }
        return aggregationBuilder;
    }

}
